package cc.shacocloud.luanniao.core.service.ai.doubao

import com.volcengine.ark.runtime.model.completion.chat.*
import com.volcengine.ark.runtime.service.ArkService
import org.springframework.ai.chat.model.ChatModel
import org.springframework.ai.chat.model.ChatResponse
import org.springframework.ai.chat.prompt.ChatOptions
import org.springframework.ai.chat.prompt.Prompt
import reactor.core.publisher.Flux


/**
 * 火山方舟的豆包大模型
 *
 * https://www.volcengine.com/docs/82379/1302010
 * @author 思追(shaco)
 */
class DoubaoChatModel(
    val arkService: ArkService,
    val endpointId: String,
    val defaultOptions: DoubaoChatOptions,
) : ChatModel, AutoCloseable {

    constructor(
        apiKey: String,
        baseUrl: String = "https://ark.cn-beijing.volces.com/api/v3/",
        endpointId: String,
        defaultOptions: DoubaoChatOptions,
    ) : this(
        ArkService.builder()
            .apiKey(apiKey)
            .baseUrl(baseUrl)
            .build(),
        endpointId,
        defaultOptions
    )

    constructor(options: DoubaoAiOptions) : this(
        options.apiKey,
        options.baseUrl,
        options.endpointId,
        options.defaultOptions
    )

    override fun call(prompt: Prompt): ChatResponse {

        val messages = prompt.toMessages()

        val chatCompletionRequest = ChatCompletionRequest.builder()
            .model(endpointId)
            .messages(messages)
            .build()

        val result = arkService.createChatCompletion(chatCompletionRequest)
        return result.toChatResponse()
    }

    override fun stream(prompt: Prompt): Flux<ChatResponse> {
        val messages = prompt.toMessages()

        val chatCompletionRequest = ChatCompletionRequest.builder()
            .model(endpointId)
            .messages(messages)
            .build()

        return Flux.from(arkService.streamChatCompletion(chatCompletionRequest))
            .map {
                it.toChatResponse()
            }
    }

    override fun getDefaultOptions(): ChatOptions {
        return defaultOptions
    }

    override fun close() {
        arkService.shutdownExecutor()
    }


}
